package com.liunian.protection;

/**
 * 基于令牌桶算法的限流器
 */
public class TokenBuketRateLimiter implements RateLimiter {

    // 令牌桶是个什么？令牌是什么？
    // 代表令牌的数量，>0说明有令牌，能放行，放行就减一，==0，无令牌 阻拦
    private int tokens;
    // 限流的本质就是令牌数
    private final int capacity;
    // 令牌桶的令牌，如果没了要怎么办？ 按照一定的速率给令牌桶加令牌，如每秒加500个，不能超过总数
    // 可以用定时任务去加 --> 启动一个定时任务，每秒执行一次 tokens + 500 不能超过capacity
    // 对于单机版的限流器可以有更简单的操作，每一次有请求要发送的时候给他加一下就好了。
    private final int rate;

    // 上一次放令牌的时间
    private Long lastTokenTime;

    public TokenBuketRateLimiter(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.lastTokenTime = System.currentTimeMillis();
        this.tokens = capacity;
    }

    /**
     * 判断请求是否可以放行
     * @return true 放行   false 拦截
     */
    public synchronized boolean allowRequest() {
        // 1、给令牌桶添加令牌
        // 计算从现在到上一次的时间间隔需要添加的令牌数
        Long currentTime = System.currentTimeMillis();
        Long timeInterval = currentTime - lastTokenTime;
        // 如果间隔时间超过500毫秒，放令牌
        if (timeInterval >= 500) {
            int needAndTokens = (int) (timeInterval * rate / 1000);
            // 给令牌桶添加令牌
            tokens = Math.min(capacity, tokens + needAndTokens);
            // 标记最后一次放入令牌的时间。
            this.lastTokenTime = System.currentTimeMillis();
        }
        // 2、自己获取令牌，如果令牌桶中有令牌则放行，否则拦截。
        if (tokens > 0) {
            tokens--;
            System.out.println("请求被放行-----------");
            return true;
        } else {
            System.out.println("请求被拦截-----------");
            return false;
        }
    }

    public static void main(String[] args) {
        TokenBuketRateLimiter tokenBuketRateLimiter = new TokenBuketRateLimiter(10, 10);
        for (int i = 0; i < 1000; i++) {
            try {
                Thread.sleep(10);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            boolean allowRequest = tokenBuketRateLimiter.allowRequest();
            System.out.println("allowRequest = " + allowRequest);
        }
    }

}
